import sys

from PyQt6.QtWidgets import QApplication, QWidget, QHBoxLayout, QVBoxLayout, \
    QTreeWidget, QTreeWidgetItem, QPushButton, QLineEdit


class MainWindow(QWidget):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle('TreeWidget 测试')
        # 命令按钮
        pbLayout = QHBoxLayout()
        pbAdd = QPushButton('添加')
        pbAdd.clicked.connect(self.addTreeNode)
        pbEdit = QPushButton('修改')
        pbEdit.clicked.connect(self.editTreeNode)
        pbDel = QPushButton('删除')
        pbDel.clicked.connect(self.delTreeNode)
        pbLayout.addWidget(pbAdd)
        pbLayout.addWidget(pbEdit)
        pbLayout.addWidget(pbDel)
        # tree
        self.tree = QTreeWidget(self)
        self.tree.setColumnCount(1)
        self.tree.setHeaderLabels(['名称'])
        root = QTreeWidgetItem(self.tree)
        root.setText(0, '学校')
        child1 = QTreeWidgetItem(root, ['学院 1'])
        child2 = QTreeWidgetItem(root, ['学院 2'])
        child21 = QTreeWidgetItem(child2, ['系 1'])
        self.tree.addTopLevelItem(root)
        self.tree.expandAll()

        self.le = QLineEdit()
        layout = QVBoxLayout()
        layout.addWidget(self.le)
        layout.addLayout(pbLayout)
        layout.addWidget(self.tree)
        self.setLayout(layout)

    def addTreeNode(self):
        item = self.tree.currentItem()
        child = QTreeWidgetItem(item)
        txt = self.le.text()
        if txt != '':
            child.setText(0, txt)

    def editTreeNode(self):
        item = self.tree.currentItem()
        txt = self.le.text()
        if txt != '':
            item.setText(0, txt)

    def delTreeNode(self):
        root = self.tree.invisibleRootItem()
        for item in self.tree.selectedItems():
            (item.parent() or root).removeChild(item)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
